{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import math\n",
    "\n",
    "import mindspore as ms\n",
    "import mindspore.nn as nn\n",
    "import mindspore.numpy as mnp\n",
    "from mindspore import ops\n",
    "from mindspore import Parameter, Tensor\n",
    "from mindspore.dataset import text\n",
    "from mindspore.common import dtype as mstype\n",
    "from mindspore.common.initializer import Uniform, HeUniform, initializer\n",
    "\n",
    "import mindnlp\n",
    "from mindnlp._legacy.abc import Seq2vecModel\n",
    "from mindnlp.modules import Glove, StaticLSTM\n",
    "from mindnlp.transforms import BasicTokenizer\n",
    "\n",
    "from tqdm import tqdm\n",
    "from bidaf.evaluate import evaluate\n",
    "from bidaf.squad_process import SQuAD1_Process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['id', 'context', 'question', 'answers', 'answer_start']\n"
     ]
    }
   ],
   "source": [
    "# load datasets\n",
    "squad_train, squad_dev = mindnlp.load_dataset('squad1')\n",
    "print(squad_train.get_col_names())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load vocab and embedding\n",
    "char_dic = {\"<unk>\": 0, \"<pad>\": 1, \"e\": 2, \"t\": 3, \"a\": 4, \"i\": 5, \"n\": 6,\\\n",
    "                    \"o\": 7, \"s\": 8, \"r\": 9, \"h\": 10, \"l\": 11, \"d\": 12, \"c\": 13, \"u\": 14,\\\n",
    "                    \"m\": 15, \"f\": 16, \"p\": 17, \"g\": 18, \"w\": 19, \"y\": 20, \"b\": 21, \",\": 22,\\\n",
    "                    \"v\": 23, \".\": 24, \"k\": 25, \"1\": 26, \"0\": 27, \"x\": 28, \"2\": 29, \"\\\"\": 30, \\\n",
    "                    \"-\": 31, \"j\": 32, \"9\": 33, \"'\": 34, \")\": 35, \"(\": 36, \"?\": 37, \"z\": 38,\\\n",
    "                    \"5\": 39, \"8\": 40, \"q\": 41, \"3\": 42, \"4\": 43, \"7\": 44, \"6\": 45, \";\": 46,\\\n",
    "                    \":\": 47, \"\\u2013\": 48, \"%\": 49, \"/\": 50, \"]\": 51, \"[\": 52}\n",
    "char_vocab = text.Vocab.from_dict(char_dic)\n",
    "# you can download the vocab file from \"https://download.mindspore.cn/toolkits/mindnlp/vocab/Glove/glove.6B.100d.txt\"\n",
    "word_vocab = text.Vocab.from_file(\"glove.6B.100d.vocab.txt\", special_tokens=[\"<unk>\", \"<pad>\"], special_first=True)\n",
    "word_embeddings = Glove.from_pretrained('6B', 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=============ready to process dataset===========\n",
      "===================process over=================\n"
     ]
    }
   ],
   "source": [
    "# process dataset\n",
    "tokenizer = BasicTokenizer(True)\n",
    "\n",
    "print(\"=============ready to process dataset===========\")\n",
    "squad_train = SQuAD1_Process(squad_train, char_vocab, word_vocab, tokenizer=tokenizer,\\\n",
    "                   max_context_len=768, max_question_len=64, max_char_len=48,\\\n",
    "                   batch_size=8, drop_remainder=False )\n",
    "squad_dev = SQuAD1_Process(squad_dev, char_vocab, word_vocab, tokenizer=tokenizer,\\\n",
    "                   max_context_len=768, max_question_len=64, max_char_len=48,\\\n",
    "                   batch_size=8, drop_remainder=False )\n",
    "print(\"===================process over=================\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct bidirectional attention flow model(BiDAF)\n",
    "def arange(start, stop, step, dtype):\n",
    "    return Tensor(mnp.arange(start, stop, step), dtype)\n",
    "\n",
    "def sequence_mask(lengths, maxlen):\n",
    "    \"\"\"generate mask matrix by seq_length\"\"\"\n",
    "    range_vector = arange(0, maxlen, 1, lengths.dtype)\n",
    "    result = range_vector < lengths.view(lengths.shape + (1,))\n",
    "    result = result.transpose((1, 0))\n",
    "    return result.astype(lengths.dtype)\n",
    "\n",
    "def select_by_mask(inputs, mask):\n",
    "    \"\"\"mask hiddens by mask matrix\"\"\"\n",
    "    return mask.view(mask.shape + (1,)).swapaxes(0, 1) \\\n",
    "        .expand_as(inputs).astype(mstype.bool_)  * inputs\n",
    "\n",
    "def get_hidden(output, seq_length):\n",
    "    \"\"\"get hidden state by seq_length\"\"\"\n",
    "    batch_index = arange(0, seq_length.shape[0], 1, seq_length.dtype)\n",
    "    indices = ops.concat((seq_length.view(-1, 1) - 1, batch_index.view(-1, 1)), 1)\n",
    "    return ops.gather_nd(output, indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build encoder\n",
    "class Encoder(nn.Module):\n",
    "    \"\"\"\n",
    "    Encoder for BiDAF model\n",
    "    \"\"\"\n",
    "    def __init__(self, char_vocab_size, char_vocab, char_dim, char_channel_size, char_channel_width, word_vocab,\n",
    "                  word_embeddings, hidden_size, dropout):\n",
    "        super().__init__()\n",
    "        self.char_vocab = char_vocab\n",
    "        self.char_dim = char_dim\n",
    "        self.char_channel_width = char_channel_width\n",
    "        self.char_channel_size = char_channel_size\n",
    "        self.word_vocab = word_vocab\n",
    "        self.hidden_size = hidden_size\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        self.init_embed = initializer(Uniform(0.001), [char_vocab_size, char_dim])\n",
    "        self.embed = Parameter(self.init_embed, name='embed')\n",
    "\n",
    "        # 1. Character Embedding Layer\n",
    "        self.char_emb = Glove(init_embed=self.embed, dropout=0.0)\n",
    "        self.char_conv = nn.Sequential(\n",
    "            nn.Conv2d(1, char_channel_size, (char_dim, char_channel_width), pad_mode=\"pad\",\n",
    "                      weight_init=HeUniform(math.sqrt(5)), bias_init=Uniform(1 / math.sqrt(1))),\n",
    "            nn.ReLU()\n",
    "            )\n",
    "\n",
    "        # 2. Word Embedding Layer\n",
    "        self.word_emb = word_embeddings\n",
    "\n",
    "        # highway network\n",
    "        self.highway_linear0 = nn.Dense(hidden_size * 2, hidden_size * 2,\n",
    "                                        weight_init=HeUniform(math.sqrt(5)),\n",
    "                                        bias_init=Uniform(1 / math.sqrt(hidden_size * 2)),\n",
    "                                        activation=nn.ReLU())\n",
    "        self.highway_linear1 = nn.Dense(hidden_size * 2, hidden_size * 2,\n",
    "                                        weight_init=HeUniform(math.sqrt(5)),\n",
    "                                        bias_init=Uniform(1 / math.sqrt(hidden_size * 2)),\n",
    "                                        activation=nn.ReLU())\n",
    "        self.highway_gate0 = nn.Dense(hidden_size * 2, hidden_size * 2,\n",
    "                                      weight_init=HeUniform(math.sqrt(5)),\n",
    "                                      bias_init=Uniform(1 / math.sqrt(hidden_size * 2)),\n",
    "                                      activation=nn.Sigmoid())\n",
    "        self.highway_gate1 = nn.Dense(hidden_size * 2, hidden_size * 2,\n",
    "                                      weight_init=HeUniform(math.sqrt(5)),\n",
    "                                      bias_init=Uniform(1 / math.sqrt(hidden_size * 2)),\n",
    "                                      activation=nn.Sigmoid())\n",
    "\n",
    "        # 3. Contextual Embedding Layer\n",
    "        self.context_LSTM = StaticLSTM(input_size=hidden_size * 2, hidden_size=hidden_size,\n",
    "                                    bidirectional=True, batch_first=True, dropout=dropout)\n",
    "    \n",
    "    def construct(self, c_char, q_char, c_word, q_word, c_lens, q_lens):\n",
    "        # 1. Character Embedding Layer\n",
    "        c_char = self.char_emb_layer(c_char)\n",
    "        q_char = self.char_emb_layer(q_char)\n",
    "\n",
    "        # 2. Word Embedding Layer\n",
    "        c_word = self.word_emb(c_word)\n",
    "        q_word = self.word_emb(q_word)\n",
    "\n",
    "        # Highway network\n",
    "        c = self.highway_network(c_char, c_word)\n",
    "        q = self.highway_network(q_char, q_word)\n",
    "        \n",
    "        # 3. Contextual Embedding Layer\n",
    "        c, _ = self.context_LSTM(c)\n",
    "        mask = sequence_mask(c_lens, c.shape[1])\n",
    "        c = select_by_mask(c, mask)\n",
    "\n",
    "        q, _ = self.context_LSTM(q)\n",
    "        mask = sequence_mask(q_lens, q.shape[1])\n",
    "        q = select_by_mask(q, mask)\n",
    "\n",
    "        return c, q\n",
    "\n",
    "    def char_emb_layer(self, x):\n",
    "        \"\"\"\n",
    "        param x: (batch, seq_len, word_len)\n",
    "        return: (batch, seq_len, char_channel_size)\n",
    "        \"\"\"\n",
    "        batch_size = x.shape[0]\n",
    "        # x: [batch, seq_len, word_len, char_dim]\n",
    "        x = self.dropout(self.char_emb(x))\n",
    "        # x: [batch, seq_len, char_dim, word_len]\n",
    "        x = ops.transpose(x, (0, 1, 3, 2))\n",
    "        # x: [batch * seq_len, 1, char_dim, word_len]\n",
    "        x = x.view(-1, self.char_dim, x.shape[3]).expand_dims(1)\n",
    "        # x: [batch * seq_len, char_channel_size, 1, conv_len] -> [batch * seq_len, char_channel_size, conv_len]\n",
    "        x = self.char_conv(x).squeeze(2)\n",
    "        # x: [batch * seq_len, char_channel_size]\n",
    "        x = ops.max(x, axis=2)[1]\n",
    "        # x: [batch, seq_len, char_channel_size]\n",
    "        x = x.view(batch_size, -1, self.char_channel_size)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def highway_network(self, x1, x2):\n",
    "        \"\"\"\n",
    "        param x1: (batch, seq_len, char_channel_size)\n",
    "        param x2: (batch, seq_len, word_dim)\n",
    "        return: (batch, seq_len, hidden_size * 2)\n",
    "        \"\"\"\n",
    "        # [batch, seq_len, char_channel_size + word_dim]\n",
    "        x = ops.concat((x1, x2), axis=-1)\n",
    "        h = self.highway_linear0(x)\n",
    "        g = self.highway_gate0(x)\n",
    "        x = g * h + (1 - g) * x\n",
    "        h = self.highway_linear1(x)\n",
    "        g = self.highway_gate1(x)\n",
    "        x = g * h + (1 - g) * x\n",
    "\n",
    "        # [batch, seq_len, hidden_size * 2]\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build head\n",
    "class Head(nn.Module):\n",
    "    \"\"\"\n",
    "    Head for BiDAF model\n",
    "    \"\"\"\n",
    "    def __init__(self, hidden_size, dropout):\n",
    "        super().__init__()\n",
    "        # 4. Attention Flow Layer\n",
    "        self.att_weight_c = nn.Dense(hidden_size * 2, 1,\n",
    "                                     weight_init=HeUniform(math.sqrt(5)),\n",
    "                                     bias_init=Uniform(1 / math.sqrt(hidden_size * 2)))\n",
    "        self.att_weight_q = nn.Dense(hidden_size * 2, 1,\n",
    "                                     weight_init=HeUniform(math.sqrt(5)),\n",
    "                                     bias_init=Uniform(1 / math.sqrt(hidden_size * 2)))\n",
    "        self.att_weight_cq = nn.Dense(hidden_size * 2, 1,\n",
    "                                      weight_init=HeUniform(math.sqrt(5)),\n",
    "                                      bias_init=Uniform(1 / math.sqrt(hidden_size * 2)))\n",
    "        self.softmax = nn.Softmax(axis=-1)\n",
    "        self.batch_matmul = ops.BatchMatMul()\n",
    "\n",
    "        # 5. Modeling Layer\n",
    "        self.modeling_LSTM1 = StaticLSTM(input_size=hidden_size * 8, hidden_size=hidden_size,\n",
    "                                      bidirectional=True, batch_first=True, dropout=dropout)\n",
    "        self.modeling_LSTM2 = StaticLSTM(input_size=hidden_size * 2, hidden_size=hidden_size,\n",
    "                                      bidirectional=True, batch_first=True, dropout=dropout)\n",
    "        \n",
    "        # 6. Output Layer\n",
    "        self.p1_weight_g = nn.Dense(hidden_size * 8, 1,\n",
    "                                    weight_init=HeUniform(math.sqrt(5)),\n",
    "                                    bias_init=Uniform(1 / math.sqrt(hidden_size * 8)))\n",
    "        self.p1_weight_m = nn.Dense(hidden_size * 2, 1,\n",
    "                                    weight_init=HeUniform(math.sqrt(5)),\n",
    "                                    bias_init=Uniform(1 / math.sqrt(hidden_size * 2)))\n",
    "        self.p2_weight_g = nn.Dense(hidden_size * 8, 1,\n",
    "                                    weight_init=HeUniform(math.sqrt(5)),\n",
    "                                    bias_init=Uniform(1 / math.sqrt(hidden_size * 8)))\n",
    "        self.p2_weight_m = nn.Dense(hidden_size * 2, 1,\n",
    "                                    weight_init=HeUniform(math.sqrt(5)),\n",
    "                                    bias_init=Uniform(1 / math.sqrt(hidden_size * 2)))\n",
    "\n",
    "        self.output_LSTM = StaticLSTM(input_size=hidden_size * 2, hidden_size=hidden_size,\n",
    "                                   bidirectional=True, batch_first=True, dropout=dropout)\n",
    "\n",
    "    def construct(self, c, q, c_lens):\n",
    "        # 4. Attention Flow Layer\n",
    "        g = self.att_flow_layer(c, q)  #c, q are generated from Contextual Embedding Layer in Encoder\n",
    "        \n",
    "        # 5. Modeling Layer\n",
    "        m, _ = self.modeling_LSTM1(g)\n",
    "        mask = sequence_mask(c_lens, g.shape[1])\n",
    "        m = select_by_mask(m, mask)\n",
    "\n",
    "        m, _ = self.modeling_LSTM2(m)\n",
    "        mask = sequence_mask(c_lens, m.shape[1])\n",
    "        m = select_by_mask(m, mask)\n",
    "\n",
    "        # 6. Output Layer\n",
    "        p1, p2 = self.output_layer(g, m, c_lens)\n",
    "\n",
    "        # [batch, c_len], [batch, c_len]\n",
    "        return p1, p2\n",
    "\n",
    "    def att_flow_layer(self, c, q):\n",
    "        \"\"\"\n",
    "        param c: (batch, c_len, hidden_size * 2)\n",
    "        param q: (batch, q_len, hidden_size * 2)\n",
    "        return: (batch, c_len, q_len)\n",
    "        \"\"\"\n",
    "        c_len = c.shape[1]\n",
    "        q_len = q.shape[1]\n",
    "\n",
    "        cq = []\n",
    "        for i in range(q_len):\n",
    "            # qi: [batch, 1, hidden_size * 2]\n",
    "            qi = q.gather(ms.Tensor(i), axis=1).expand_dims(1)\n",
    "            # ci: [batch, c_len, 1] -> [batch, c_len]\n",
    "            ci = self.att_weight_cq(c * qi).squeeze(2)\n",
    "            cq.append(ci)\n",
    "        # cq: [batch, c_len, q_len]\n",
    "        cq = ops.stack(cq, -1)\n",
    "\n",
    "        # s: [batch, c_len, q_len]\n",
    "        s = self.att_weight_c(c).broadcast_to((-1, -1, q_len)) + \\\n",
    "            self.att_weight_q(q).transpose((0, 2, 1)).broadcast_to((-1, c_len, -1)) + cq\n",
    "\n",
    "        # a: [batch, c_len, q_len]\n",
    "        a = self.softmax(s)\n",
    "        # c2q_att: [batch, c_len, hidden_size * 2]\n",
    "        c2q_att = self.batch_matmul(a, q)\n",
    "        # b: [batch, 1, c_len]\n",
    "        b = self.softmax(ops.max(s, axis=2)[1]).expand_dims(1)\n",
    "        # q2c_att: [batch, hidden_size * 2]\n",
    "        q2c_att = self.batch_matmul(b, c).squeeze(1)\n",
    "        # q2c_att: [batch, c_len, hidden_size * 2]\n",
    "        q2c_att = q2c_att.expand_dims(1).broadcast_to((-1, c_len, -1))\n",
    "\n",
    "        # x: [batch, c_len, hidden_size * 8]\n",
    "        x = ops.concat([c, c2q_att, c * c2q_att, c * q2c_att], axis=-1)\n",
    "        return x\n",
    "\n",
    "    def output_layer(self, g, m, l):\n",
    "        \"\"\"\n",
    "        param g: (batch, c_len, hidden_size * 8)\n",
    "        param m: (batch, c_len ,hidden_size * 2)\n",
    "        return: p1: (batch, c_len), p2: (batch, c_len)\n",
    "        \"\"\"\n",
    "        # p1: [batch, c_len]\n",
    "        p1 = (self.p1_weight_g(g) + self.p1_weight_m(m)).squeeze(2)\n",
    "        # m2: [batch, c_len, hidden_size * 2]\n",
    "        m2, _ = self.output_LSTM(m)\n",
    "        mask = sequence_mask(l, m.shape[1])\n",
    "        m2 = select_by_mask(m2, mask)\n",
    "        # p2: [batch, c_len]\n",
    "        p2 = (self.p2_weight_g(g) + self.p2_weight_m(m2)).squeeze(2)\n",
    "\n",
    "        return p1, p2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BiDAF(Seq2vecModel):\n",
    "    def __init__(self, encoder, head):\n",
    "        super().__init__(encoder, head)\n",
    "        self.encoder = encoder\n",
    "        self.head = head\n",
    "\n",
    "    def construct(self, c_char, q_char, c_word, q_word, c_lens, q_lens):\n",
    "        c, q = self.encoder(c_char, q_char, c_word, q_word, c_lens, q_lens)\n",
    "        p1, p2 = self.head(c, q, c_lens)\n",
    "        return p1, p2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define some parameters\n",
    "char_vocab_size = len(char_vocab.vocab())\n",
    "char_dim = 8\n",
    "char_channel_width = 5\n",
    "char_channel_size = 100\n",
    "hidden_size = 100\n",
    "dropout = 0.2\n",
    "lr = 0.5\n",
    "epochs = 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(2655801:139853446842176,MainProcess):2023-04-09-01:09:08.311.591 [mindspore/nn/layer/basic.py:162] For Dropout, this parameter `keep_prob` will be deprecated, please use `p` instead.\n",
      "[WARNING] ME(2655801:139853446842176,MainProcess):2023-04-09-01:09:08.316.591 [mindspore/nn/layer/conv.py:104] Value of 'has_bias' is False, value of 'bias_init' will be ignored.\n",
      "[WARNING] ME(2655801:139853446842176,MainProcess):2023-04-09-01:09:08.327.933 [/home/daiyuxin/ytt/mindnlp/mindnlp/modules/rnns.py:365] dropout option adds dropout after all but last recurrent layer, so non-zero dropout expects num_layers greater than 1, but got dropout=0.2 and num_layers=1\n",
      "[WARNING] ME(2655801:139853446842176,MainProcess):2023-04-09-01:09:08.340.816 [/home/daiyuxin/ytt/mindnlp/mindnlp/modules/rnns.py:365] dropout option adds dropout after all but last recurrent layer, so non-zero dropout expects num_layers greater than 1, but got dropout=0.2 and num_layers=1\n",
      "[WARNING] ME(2655801:139853446842176,MainProcess):2023-04-09-01:09:08.350.739 [/home/daiyuxin/ytt/mindnlp/mindnlp/modules/rnns.py:365] dropout option adds dropout after all but last recurrent layer, so non-zero dropout expects num_layers greater than 1, but got dropout=0.2 and num_layers=1\n",
      "[WARNING] ME(2655801:139853446842176,MainProcess):2023-04-09-01:09:08.364.022 [/home/daiyuxin/ytt/mindnlp/mindnlp/modules/rnns.py:365] dropout option adds dropout after all but last recurrent layer, so non-zero dropout expects num_layers greater than 1, but got dropout=0.2 and num_layers=1\n"
     ]
    }
   ],
   "source": [
    "# net\n",
    "encoder = Encoder(char_vocab_size, char_vocab, char_dim, char_channel_size, char_channel_width, word_vocab,\n",
    "                  word_embeddings, hidden_size, dropout)                  \n",
    "head = Head(hidden_size, dropout)\n",
    "net = BiDAF(encoder, head)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(2655801:139853446842176,MainProcess):2023-04-09-01:09:11.660.734 [mindspore/nn/layer/basic.py:188] For Dropout, this parameter `keep_prob` will be deprecated, please use `p` instead.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BiDAF<\n",
       "  (encoder): Encoder<\n",
       "    (dropout): Dropout<keep_prob=0.8>\n",
       "    (char_emb): Glove<\n",
       "      (dropout_layer): Dropout<p=0.0>\n",
       "      >\n",
       "    (char_conv): Sequential<\n",
       "      (0): Conv2d<input_channels=1, output_channels=100, kernel_size=(8, 5), stride=(1, 1), pad_mode=pad, padding=0, dilation=(1, 1), group=1, has_bias=False, weight_init=<mindspore.common.initializer.HeUniform object at 0x7f30d4acd690>, bias_init=<mindspore.common.initializer.Uniform object at 0x7f303f9cfe10>, format=NCHW>\n",
       "      (1): ReLU<>\n",
       "      >\n",
       "    (word_emb): Glove<\n",
       "      (dropout_layer): Dropout<p=0.0>\n",
       "      >\n",
       "    (highway_linear0): Dense<\n",
       "      input_channels=200, output_channels=200, has_bias=True, activation=ReLU<>\n",
       "      (activation): ReLU<>\n",
       "      >\n",
       "    (highway_linear1): Dense<\n",
       "      input_channels=200, output_channels=200, has_bias=True, activation=ReLU<>\n",
       "      (activation): ReLU<>\n",
       "      >\n",
       "    (highway_gate0): Dense<\n",
       "      input_channels=200, output_channels=200, has_bias=True, activation=Sigmoid<>\n",
       "      (activation): Sigmoid<>\n",
       "      >\n",
       "    (highway_gate1): Dense<\n",
       "      input_channels=200, output_channels=200, has_bias=True, activation=Sigmoid<>\n",
       "      (activation): Sigmoid<>\n",
       "      >\n",
       "    (context_LSTM): StaticLSTM<\n",
       "      (rnn): MultiLayerRNN<\n",
       "        (cell_list): CellList<\n",
       "          (0): SingleLSTMLayer_GPU<>\n",
       "          >\n",
       "        (dropout): Dropout<p=0.2>\n",
       "        >\n",
       "      >\n",
       "    >\n",
       "  (head): Head<\n",
       "    (att_weight_c): Dense<input_channels=200, output_channels=1, has_bias=True>\n",
       "    (att_weight_q): Dense<input_channels=200, output_channels=1, has_bias=True>\n",
       "    (att_weight_cq): Dense<input_channels=200, output_channels=1, has_bias=True>\n",
       "    (softmax): Softmax<>\n",
       "    (modeling_LSTM1): StaticLSTM<\n",
       "      (rnn): MultiLayerRNN<\n",
       "        (cell_list): CellList<\n",
       "          (0): SingleLSTMLayer_GPU<>\n",
       "          >\n",
       "        (dropout): Dropout<p=0.2>\n",
       "        >\n",
       "      >\n",
       "    (modeling_LSTM2): StaticLSTM<\n",
       "      (rnn): MultiLayerRNN<\n",
       "        (cell_list): CellList<\n",
       "          (0): SingleLSTMLayer_GPU<>\n",
       "          >\n",
       "        (dropout): Dropout<p=0.2>\n",
       "        >\n",
       "      >\n",
       "    (p1_weight_g): Dense<input_channels=800, output_channels=1, has_bias=True>\n",
       "    (p1_weight_m): Dense<input_channels=200, output_channels=1, has_bias=True>\n",
       "    (p2_weight_g): Dense<input_channels=800, output_channels=1, has_bias=True>\n",
       "    (p2_weight_m): Dense<input_channels=200, output_channels=1, has_bias=True>\n",
       "    (output_LSTM): StaticLSTM<\n",
       "      (rnn): MultiLayerRNN<\n",
       "        (cell_list): CellList<\n",
       "          (0): SingleLSTMLayer_GPU<>\n",
       "          >\n",
       "        (dropout): Dropout<p=0.2>\n",
       "        >\n",
       "      >\n",
       "    >\n",
       "  >"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define Loss & Optimizer\n",
    "class Loss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def construct(self, logit1, logit2, s_idx, e_idx):\n",
    "        loss_fn = nn.CrossEntropyLoss()\n",
    "        loss = loss_fn(logit1, s_idx) + loss_fn(logit2, e_idx)\n",
    "        return loss\n",
    "\n",
    "loss_fn = Loss()\n",
    "optimizer = nn.Adadelta(net.trainable_params(), learning_rate=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward_fn(c_char, q_char, c_word, q_word, c_lens, q_lens, s_idx, e_idx):\n",
    "    logits = net(c_char, q_char, c_word, q_word, c_lens, q_lens)\n",
    "    loss = loss_fn(*logits, s_idx, e_idx)\n",
    "    return_list = (loss,) + logits\n",
    "    return return_list\n",
    "\n",
    "grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)\n",
    "\n",
    "def train_step(c_char, q_char, c_word, q_word, c_lens, q_lens, s_idx, e_idx):\n",
    "    (loss, *_), grads = grad_fn(c_char, q_char, c_word, q_word, c_lens, q_lens, s_idx, e_idx)\n",
    "    optimizer(grads)\n",
    "    return loss\n",
    "\n",
    "def train_one_epoch(model, train_dataset, epoch=0):\n",
    "    model.set_train()\n",
    "    total = train_dataset.get_dataset_size()\n",
    "    loss_total = 0\n",
    "    step_total = 0\n",
    "    with tqdm(total=total) as t:\n",
    "        t.set_description('Epoch %i' % epoch)\n",
    "        for _, c_word, q_word, c_char, q_char, c_lens, q_lens, s_idx, e_idx in train_dataset.create_tuple_iterator():\n",
    "            loss = train_step(c_char, q_char, c_word, q_word, c_lens, q_lens, s_idx, e_idx)\n",
    "            loss_total += loss.asnumpy()\n",
    "            step_total += 1\n",
    "            t.set_postfix(loss=loss_total/step_total)\n",
    "            t.update(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_loop(model, dataset, vocab, loss_fn):\n",
    "    model.set_train(False)\n",
    "    loss = 0\n",
    "    answers = dict()\n",
    "\n",
    "    for ids, c_word, q_word, c_char, q_char, c_lens, q_lens, s_idx, e_idx in dataset.create_tuple_iterator():\n",
    "        p1, p2 = model(c_char, q_char, c_word, q_word, c_lens, q_lens)\n",
    "        batch_loss = loss_fn(p1, p2, s_idx, e_idx)\n",
    "        loss += batch_loss\n",
    "\n",
    "        # [batch, c_len]\n",
    "        batch_size, c_len = p1.shape\n",
    "        ls = nn.LogSoftmax(axis=1)\n",
    "        mask = mnp.tril((ops.ones((c_len, c_len), dtype=ms.float32) * float('-inf')),\n",
    "                         k=-1).expand_dims(0).broadcast_to((batch_size, -1, -1))\n",
    "        mask = mnp.where(ops.isnan(mask), ops.zeros_like(mask), mask)\n",
    "        score = (ls(p1).expand_dims(2) + ls(p2).expand_dims(1)) + mask\n",
    "        s_idx, score = ops.max(score, axis=1)\n",
    "        e_idx, score = ops.max(score, axis=1)\n",
    "        s_idx = ops.gather_elements(s_idx, 1, e_idx.view(-1, 1)).squeeze(axis=1)\n",
    "\n",
    "        for i in range(batch_size):\n",
    "            answer_id = ids.asnumpy()[i]\n",
    "            answer = c_word[i][s_idx[i].asnumpy().item():e_idx[i].asnumpy().item()+1]\n",
    "\n",
    "            answer_list = []\n",
    "            for idx in answer:\n",
    "                idx = idx.asnumpy().item()\n",
    "                if idx < 0:\n",
    "                    idx = idx + 188744\n",
    "                answer_list.append(vocab.ids_to_tokens(idx))\n",
    "            answer = ' '.join(answer_list)\n",
    "            answers[answer_id.item()] = answer\n",
    "    # you can download the squad dev dataset from \"https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json\"\n",
    "    with open(\"dev-v1.1.json\") as dataset_file:\n",
    "        dataset_json = json.load(dataset_file)\n",
    "        squad_data = dataset_json['data']\n",
    "    exact_match, f1 = evaluate(squad_data, answers)\n",
    "    print(f\"Test: \\n EM: {exact_match:.3f}, F1: {f1:.3f}, Avg loss: {loss.asnumpy().item():>8f} \\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0:   0%|          | 0/10950 [00:00<?, ?it/s][WARNING] PYNATIVE(2655801,7f322b050740,python):2023-04-09-10:11:11.927.124 [mindspore/ccsrc/pipeline/pynative/grad/grad.cc:1252] CheckAlreadyRun] The input info of this cell has changed, forward process will run again\n",
      "Epoch 0:  15%|█▌        | 1688/10950 [10:15<56:15,  2.74it/s, loss=9.61] \n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "\u001B[0;32m/tmp/ipykernel_2655801/1869801482.py\u001B[0m in \u001B[0;36m<module>\u001B[0;34m\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[0;32mfor\u001B[0m \u001B[0mepoch\u001B[0m \u001B[0;32min\u001B[0m \u001B[0mrange\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mepochs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m----> 2\u001B[0;31m     \u001B[0mtrain_one_epoch\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mnet\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0msquad_train\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mepoch\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m      3\u001B[0m     \u001B[0mtest_loop\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mnet\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0msquad_dev\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mword_vocab\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mloss_fn\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m      4\u001B[0m \u001B[0mprint\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m\"Done!\"\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/tmp/ipykernel_2655801/3207938784.py\u001B[0m in \u001B[0;36mtrain_one_epoch\u001B[0;34m(model, train_dataset, epoch)\u001B[0m\n\u001B[1;32m     20\u001B[0m         \u001B[0mt\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mset_description\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m'Epoch %i'\u001B[0m \u001B[0;34m%\u001B[0m \u001B[0mepoch\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     21\u001B[0m         \u001B[0;32mfor\u001B[0m \u001B[0m_\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0ms_idx\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0me_idx\u001B[0m \u001B[0;32min\u001B[0m \u001B[0mtrain_dataset\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mcreate_tuple_iterator\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m---> 22\u001B[0;31m             \u001B[0mloss\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mtrain_step\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mc_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0ms_idx\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0me_idx\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m     23\u001B[0m             \u001B[0mloss_total\u001B[0m \u001B[0;34m+=\u001B[0m \u001B[0mloss\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0masnumpy\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     24\u001B[0m             \u001B[0mstep_total\u001B[0m \u001B[0;34m+=\u001B[0m \u001B[0;36m1\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/tmp/ipykernel_2655801/3207938784.py\u001B[0m in \u001B[0;36mtrain_step\u001B[0;34m(c_char, q_char, c_word, q_word, c_lens, q_lens, s_idx, e_idx)\u001B[0m\n\u001B[1;32m      8\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m      9\u001B[0m \u001B[0;32mdef\u001B[0m \u001B[0mtrain_step\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mc_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0ms_idx\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0me_idx\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m---> 10\u001B[0;31m     \u001B[0;34m(\u001B[0m\u001B[0mloss\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0m_\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mgrads\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mgrad_fn\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mc_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0ms_idx\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0me_idx\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m     11\u001B[0m     \u001B[0moptimizer\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mgrads\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     12\u001B[0m     \u001B[0;32mreturn\u001B[0m \u001B[0mloss\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/ops/composite/base.py\u001B[0m in \u001B[0;36mafter_grad\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    602\u001B[0m                 \u001B[0;32mif\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mget_by_list\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    603\u001B[0m                     \u001B[0;32mdef\u001B[0m \u001B[0mafter_grad\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 604\u001B[0;31m                         \u001B[0;32mreturn\u001B[0m \u001B[0mgrad_\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfn_\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mweights\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    605\u001B[0m                 \u001B[0;32melse\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    606\u001B[0m                     \u001B[0;32mdef\u001B[0m \u001B[0mafter_grad\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/common/api.py\u001B[0m in \u001B[0;36mwrapper\u001B[0;34m(*arg, **kwargs)\u001B[0m\n\u001B[1;32m     99\u001B[0m     \u001B[0;34m@\u001B[0m\u001B[0mwraps\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfn\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    100\u001B[0m     \u001B[0;32mdef\u001B[0m \u001B[0mwrapper\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0marg\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 101\u001B[0;31m         \u001B[0mresults\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mfn\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0marg\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    102\u001B[0m         \u001B[0;32mreturn\u001B[0m \u001B[0m_convert_python_data\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mresults\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    103\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/ops/composite/base.py\u001B[0m in \u001B[0;36mafter_grad\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    579\u001B[0m             \u001B[0;34m@\u001B[0m\u001B[0m_wrap_func\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    580\u001B[0m             \u001B[0;32mdef\u001B[0m \u001B[0mafter_grad\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 581\u001B[0;31m                 \u001B[0mres\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_pynative_forward_run\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfn\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mgrad_\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mweights\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    582\u001B[0m                 \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mgrad\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfn\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mgrad_\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mweights\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mgrad_position\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    583\u001B[0m                 \u001B[0mout\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/ops/composite/base.py\u001B[0m in \u001B[0;36m_pynative_forward_run\u001B[0;34m(self, fn, grad, weights, args, kwargs)\u001B[0m\n\u001B[1;32m    627\u001B[0m                 \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mset_grad_flag\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;32mTrue\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    628\u001B[0m                 \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mnew_graph\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfn\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mnew_kwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 629\u001B[0;31m                 \u001B[0moutputs\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mfn\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mnew_kwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    630\u001B[0m                 \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mend_graph\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mfn\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0moutputs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mnew_kwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    631\u001B[0m                 \u001B[0;32mreturn\u001B[0m \u001B[0moutputs\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/ops/composite/base.py\u001B[0m in \u001B[0;36maux_fn\u001B[0;34m(*args)\u001B[0m\n\u001B[1;32m    543\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    544\u001B[0m         \u001B[0;32mdef\u001B[0m \u001B[0maux_fn\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 545\u001B[0;31m             \u001B[0moutputs\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mfn\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    546\u001B[0m             \u001B[0;32mif\u001B[0m \u001B[0;32mnot\u001B[0m \u001B[0misinstance\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0moutputs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mtuple\u001B[0m\u001B[0;34m)\u001B[0m \u001B[0;32mor\u001B[0m \u001B[0mlen\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0moutputs\u001B[0m\u001B[0;34m)\u001B[0m \u001B[0;34m<\u001B[0m \u001B[0;36m2\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    547\u001B[0m                 \u001B[0;32mraise\u001B[0m \u001B[0mValueError\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m\"When has_aux is True, origin fn requires more than one outputs.\"\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/tmp/ipykernel_2655801/3207938784.py\u001B[0m in \u001B[0;36mforward_fn\u001B[0;34m(c_char, q_char, c_word, q_word, c_lens, q_lens, s_idx, e_idx)\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[0;32mdef\u001B[0m \u001B[0mforward_fn\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mc_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0ms_idx\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0me_idx\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m----> 2\u001B[0;31m     \u001B[0mlogits\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mnet\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mc_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_lens\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m      3\u001B[0m     \u001B[0mloss\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mloss_fn\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0mlogits\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0ms_idx\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0me_idx\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m      4\u001B[0m     \u001B[0mreturn_list\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0;34m(\u001B[0m\u001B[0mloss\u001B[0m\u001B[0;34m,\u001B[0m\u001B[0;34m)\u001B[0m \u001B[0;34m+\u001B[0m \u001B[0mlogits\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m      5\u001B[0m     \u001B[0;32mreturn\u001B[0m \u001B[0mreturn_list\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/nn/cell.py\u001B[0m in \u001B[0;36m__call__\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m    651\u001B[0m         \u001B[0;32mtry\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    652\u001B[0m             \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mnew_graph\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 653\u001B[0;31m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_run_construct\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    654\u001B[0m             \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mend_graph\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0moutput\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    655\u001B[0m         \u001B[0;32mexcept\u001B[0m \u001B[0mException\u001B[0m \u001B[0;32mas\u001B[0m \u001B[0merr\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/nn/cell.py\u001B[0m in \u001B[0;36m_run_construct\u001B[0;34m(self, cast_inputs, kwargs)\u001B[0m\n\u001B[1;32m    439\u001B[0m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_shard_fn\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0mcast_inputs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    440\u001B[0m         \u001B[0;32melse\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 441\u001B[0;31m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mconstruct\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0mcast_inputs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    442\u001B[0m         \u001B[0;32mif\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_enable_forward_hook\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    443\u001B[0m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_run_forward_hook\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mcast_inputs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0moutput\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/tmp/ipykernel_2655801/2209310489.py\u001B[0m in \u001B[0;36mconstruct\u001B[0;34m(self, c_char, q_char, c_word, q_word, c_lens, q_lens)\u001B[0m\n\u001B[1;32m      7\u001B[0m     \u001B[0;32mdef\u001B[0m \u001B[0mconstruct\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_lens\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m      8\u001B[0m         \u001B[0mc\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mencoder\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mc_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_char\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_word\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_lens\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq_lens\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m----> 9\u001B[0;31m         \u001B[0mp1\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mp2\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mhead\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mc\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_lens\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m     10\u001B[0m         \u001B[0;32mreturn\u001B[0m \u001B[0mp1\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mp2\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/nn/cell.py\u001B[0m in \u001B[0;36m__call__\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m    651\u001B[0m         \u001B[0;32mtry\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    652\u001B[0m             \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mnew_graph\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 653\u001B[0;31m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_run_construct\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    654\u001B[0m             \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mend_graph\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0moutput\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    655\u001B[0m         \u001B[0;32mexcept\u001B[0m \u001B[0mException\u001B[0m \u001B[0;32mas\u001B[0m \u001B[0merr\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/nn/cell.py\u001B[0m in \u001B[0;36m_run_construct\u001B[0;34m(self, cast_inputs, kwargs)\u001B[0m\n\u001B[1;32m    439\u001B[0m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_shard_fn\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0mcast_inputs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    440\u001B[0m         \u001B[0;32melse\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 441\u001B[0;31m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mconstruct\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0mcast_inputs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    442\u001B[0m         \u001B[0;32mif\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_enable_forward_hook\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    443\u001B[0m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_run_forward_hook\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mcast_inputs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0moutput\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/tmp/ipykernel_2655801/2485060023.py\u001B[0m in \u001B[0;36mconstruct\u001B[0;34m(self, c, q, c_lens)\u001B[0m\n\u001B[1;32m     44\u001B[0m     \u001B[0;32mdef\u001B[0m \u001B[0mconstruct\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mc_lens\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     45\u001B[0m         \u001B[0;31m# 4. Attention Flow Layer\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m---> 46\u001B[0;31m         \u001B[0mg\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0matt_flow_layer\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mc\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mq\u001B[0m\u001B[0;34m)\u001B[0m  \u001B[0;31m#c, q are generated from Contextual Embedding Layer in Encoder\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m     47\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     48\u001B[0m         \u001B[0;31m# 5. Modeling Layer\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/tmp/ipykernel_2655801/2485060023.py\u001B[0m in \u001B[0;36matt_flow_layer\u001B[0;34m(self, c, q)\u001B[0m\n\u001B[1;32m     75\u001B[0m             \u001B[0mqi\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mq\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mgather\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mms\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mTensor\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mi\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0maxis\u001B[0m\u001B[0;34m=\u001B[0m\u001B[0;36m1\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mexpand_dims\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;36m1\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     76\u001B[0m             \u001B[0;31m# ci: [batch, c_len, 1] -> [batch, c_len]\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m---> 77\u001B[0;31m             \u001B[0mci\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0matt_weight_cq\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mc\u001B[0m \u001B[0;34m*\u001B[0m \u001B[0mqi\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0msqueeze\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;36m2\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m     78\u001B[0m             \u001B[0mcq\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mappend\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mci\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     79\u001B[0m         \u001B[0;31m# cq: [batch, c_len, q_len]\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/nn/cell.py\u001B[0m in \u001B[0;36m__call__\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m    651\u001B[0m         \u001B[0;32mtry\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    652\u001B[0m             \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mnew_graph\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 653\u001B[0;31m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_run_construct\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    654\u001B[0m             \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mend_graph\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0moutput\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    655\u001B[0m         \u001B[0;32mexcept\u001B[0m \u001B[0mException\u001B[0m \u001B[0;32mas\u001B[0m \u001B[0merr\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/nn/cell.py\u001B[0m in \u001B[0;36m_run_construct\u001B[0;34m(self, cast_inputs, kwargs)\u001B[0m\n\u001B[1;32m    439\u001B[0m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_shard_fn\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0mcast_inputs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    440\u001B[0m         \u001B[0;32melse\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 441\u001B[0;31m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mconstruct\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m*\u001B[0m\u001B[0mcast_inputs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    442\u001B[0m         \u001B[0;32mif\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_enable_forward_hook\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    443\u001B[0m             \u001B[0moutput\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_run_forward_hook\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mcast_inputs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0moutput\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/nn/layer/basic.py\u001B[0m in \u001B[0;36mconstruct\u001B[0;34m(self, x)\u001B[0m\n\u001B[1;32m    583\u001B[0m         \u001B[0;32mif\u001B[0m \u001B[0mlen\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mx_shape\u001B[0m\u001B[0;34m)\u001B[0m \u001B[0;34m!=\u001B[0m \u001B[0;36m2\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    584\u001B[0m             \u001B[0mout_shape\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mx_shape\u001B[0m\u001B[0;34m[\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m-\u001B[0m\u001B[0;36m1\u001B[0m\u001B[0;34m]\u001B[0m \u001B[0;34m+\u001B[0m \u001B[0;34m(\u001B[0m\u001B[0mF\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mshape\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mx\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m[\u001B[0m\u001B[0;34m-\u001B[0m\u001B[0;36m1\u001B[0m\u001B[0;34m]\u001B[0m\u001B[0;34m,\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 585\u001B[0;31m             \u001B[0mx\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mreshape\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mx\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mout_shape\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    586\u001B[0m         \u001B[0;32mreturn\u001B[0m \u001B[0mx\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    587\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/ops/primitive.py\u001B[0m in \u001B[0;36m__call__\u001B[0;34m(self, *args)\u001B[0m\n\u001B[1;32m    315\u001B[0m         \u001B[0;32mif\u001B[0m \u001B[0mshould_elim\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    316\u001B[0m             \u001B[0;32mreturn\u001B[0m \u001B[0moutput\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 317\u001B[0;31m         \u001B[0;32mreturn\u001B[0m \u001B[0m_run_op\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mname\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0margs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    318\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    319\u001B[0m     \u001B[0;32mdef\u001B[0m \u001B[0m__getstate__\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/ops/primitive.py\u001B[0m in \u001B[0;36m_run_op\u001B[0;34m(obj, op_name, args)\u001B[0m\n\u001B[1;32m    881\u001B[0m     \u001B[0;34m\"\"\"Single op execution function supported by ge in PyNative mode.\"\"\"\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    882\u001B[0m     \u001B[0;32mif\u001B[0m \u001B[0m_RUN_OP_ASYNC\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m--> 883\u001B[0;31m         \u001B[0mstub\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0m_pynative_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mrun_op_async\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mobj\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0margs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m    884\u001B[0m         \u001B[0;32mreturn\u001B[0m \u001B[0m_convert_stub\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mstub\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m    885\u001B[0m     \u001B[0;32mreturn\u001B[0m \u001B[0m_run_op_sync\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mobj\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mop_name\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0margs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m~/anaconda3/envs/ytt/lib/python3.7/site-packages/mindspore/common/api.py\u001B[0m in \u001B[0;36mrun_op_async\u001B[0;34m(self, prim, args)\u001B[0m\n\u001B[1;32m   1050\u001B[0m             \u001B[0mStubNode\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mresult\u001B[0m \u001B[0mof\u001B[0m \u001B[0mrun\u001B[0m \u001B[0mop\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m   1051\u001B[0m         \"\"\"\n\u001B[0;32m-> 1052\u001B[0;31m         \u001B[0;32mreturn\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_executor\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mrun_op_async\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mprim\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0margs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m   1053\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m   1054\u001B[0m     \u001B[0;32mdef\u001B[0m \u001B[0mnew_graph\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mself\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mobj\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m*\u001B[0m\u001B[0margs\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m**\u001B[0m\u001B[0mkwargs\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "source": [
    "for epoch in range(epochs):\n",
    "    train_one_epoch(net, squad_train, epoch)\n",
    "    test_loop(net, squad_dev, word_vocab, loss_fn)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "======================="
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ytt",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
